package learner.hmm;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Beta {

	private Matrix<Integer, Integer> a;

	private Matrix<Integer, MovementVector> b;

	private Movement movement;

	// list of already calculated alpha's, the index represents the time
	// (k-variable on slide)
	// the map is from a state to a probability.
	private List<Map<Integer, Double>> betaKs;

	public Beta(Matrix<Integer,Integer> a,
			Matrix<Integer, MovementVector> b, Movement movement) {
		this.a = a;
		this.b = b;
		this.movement = movement;
		betaKs = new ArrayList<Map<Integer, Double>>();
		initializeBeta();
	}

	private void initializeBeta() {
		for (int i = 0; i < BaumWelchLearner.NB_OF_STATES; i++) {
			double beta_K_i = 1;
			setValue(movement.getSize(), i, beta_K_i);
		}
	}

	public double get(int k, int j) {
		// trivial case:
		if (valueKnown(k, j))
			return getKnownValue(k, j);
		else { // recursion
			double sum = 0;
			for (int i = 0 ; i < BaumWelchLearner.NB_OF_STATES; i ++) {
				sum += get(k+1, i) * a.get(j,i); // a_j_i in formule
			}
			setValue(k, j, sum * b.get(j,movement.getMovementForTime(k + 1)));
			return sum * b.get(j, movement.getMovementForTime(k + 1));
		}
	}

	private boolean valueKnown(int k, int i) {
		int kComplement = movement.getSize() - k;
		if (betaKs.size() <= kComplement)
			return false;
		if (!betaKs.get(kComplement).containsKey(i))
			return false;
		return true;
	}
	
	// k can range from 1 to nbof movements
	private void setValue(int k, int i, double value) {
		int kComplement = movement.getSize() - k;
		if (kComplement > betaKs.size())
			throw new IllegalStateException(
					"You can't know the value yet: known betaKs= "
							+ betaKs.size() + " given time= " + k);
		if (valueKnown(k, i))
			throw new IllegalStateException("Beta value already known: k= " + k + " i = " + i);
		if (kComplement == betaKs.size()) {
			Map<Integer, Double> alphas_k = new HashMap<Integer, Double>();
			alphas_k.put(i, value);
			betaKs.add(alphas_k);
			
		} else {
			betaKs.get(kComplement).put(i, value);
		}
	}
	
	private double getKnownValue(int k, int i) {
		int kComplement = movement.getSize() - k;
		return betaKs.get(kComplement).get(i);
	}

}
